#include <webots/Robot.hpp>
#include <webots/Gyro.hpp>
#include <webots/GPS.hpp>
#include <webots/Motor.hpp>
#include <webots/InertialUnit.hpp>
#include <webots/LED.hpp>

#include <Eigen/Eigen>
#include <common.h>

#define MASS 2.70703                   // kg
#define G 9.8                          // m/s2
#define J11 0.0341171                  // inertial matrix Ixx
#define J22 0.0363089                  // inertial matrix Iyy
#define J33 0.0627386                  // inertial matrix Izz
#define J12 0                          // inertial matrix Ixy
#define J13 0                          // inertial matrix Ixz
#define J23 0                          // inertial matrix Iyz
#define KF           2.55e-5           // N  per rad/s
#define KM           5.1e-7            // Nm per rad/s
#define ARM_LENGTH   0.22978           // m

using namespace webots;
using namespace Eigen;

Vector3d rotation_matrix_to_eulerZYX(Matrix3d R)
{
  double r11 = R(0, 0);
  double r21 = R(1, 0);
  double r31 = R(2, 0);
  double r32 = R(2, 1);
  double r33 = R(2, 2);
  Vector3d euler;
  euler.z() = std::atan2(r21, r11);
  euler.x() = std::atan2(r32, r33);
  euler.y() = std::atan2(-r31, std::sqrt(r11 * r11 + r21 * r21));
  return euler;
};

int main(int argc, char **argv) 
{
  Robot* robot = new Robot;
  int timestep = (int)robot->getBasicTimeStep();
  Motor* motor1 = robot->getMotor("front_right_motor");
  Motor* motor2 = robot->getMotor("front_left_motor");
  Motor* motor3 = robot->getMotor("rear_left_motor");
  Motor* motor4 = robot->getMotor("rear_right_motor");
  motor1->setPosition(INFINITY);
  motor2->setPosition(INFINITY);
  motor3->setPosition(INFINITY);
  motor4->setPosition(INFINITY);
  motor1->setVelocity(0);
  motor2->setVelocity(0);
  motor3->setVelocity(0);
  motor4->setVelocity(0);
  GPS* gps = robot->getGPS("gps");
  gps->enable(timestep);
  Gyro* gyro = robot->getGyro("gyro");
  gyro->enable(timestep);
  InertialUnit* imu = robot->getInertialUnit("inertial unit");
  imu->enable(timestep);
  LED* led = robot->getLED("led");
  led->set(1);

  while (robot->step(timestep) != -1) 
  {
    // Read the sensors:
    const double* gps_data = gps->getValues();
    const double* imu_data = imu->getQuaternion();
    const double* gyro_data = gyro->getValues();
    Vector3d pos;
    pos.x() =  gps_data[0];
    pos.y() = -gps_data[2];
    pos.z() =  gps_data[1];
    static Vector3d prev_x = pos;
    Vector3d vel = (pos - prev_x) / (timestep / 1e3);
    prev_x = pos;
    Vector3d Om;
    Om.x() = gyro_data[0];
    Om.y() = gyro_data[1];
    Om.z() = gyro_data[2];
    Quaterniond q = Quaterniond(cos(M_PI / 4), sin(M_PI / 4), 0, 0) * Quaterniond(imu_data[3], imu_data[0], imu_data[1], imu_data[2]);
    Matrix3d R = q.matrix();
    Vector3d euler = rotation_matrix_to_eulerZYX(R);
    INFO_STREAM("custom_uav_control", "position: z=" << pos.z() << ", velocity: z=" << vel.z() << ", roll: " << euler.x() << ", pitch: " << euler.y() << ", yaw: " << euler.z());

    // Process sensor data here.
    Vector4d input(10, -10, 10, -10);

    // Enter here functions to send actuator commands, like:
    motor1->setVelocity(input(0));
    motor2->setVelocity(input(1));
    motor3->setVelocity(input(2));
    motor4->setVelocity(input(3));
  };

  delete robot;
  return 0;
}
